//
// Created by v on 24-12-24.
//

#include <stdlib.h>
#include <string.h>

#include "CMAIServiceImplRk3588.h"
#include "postprocess.h"
#include "rt_print.h"

CMAIServiceImplRk3588::CMAIServiceImplRk3588() {

}

CMAIServiceImplRk3588::~CMAIServiceImplRk3588() noexcept {}

int CMAIServiceImplRk3588::init() {

    return 0;
}

int CMAIServiceImplRk3588::deinit() {

    return 0;
}

int CMAIServiceImplRk3588::init_model(const char *model_path, char *label_file) {
    memset(&rknn_app_ctx, 0, sizeof(rknn_app_context_t));
    init_post_process(label_file);
    int ret = init_yolo11_model(model_path, &rknn_app_ctx);
    if (ret != 0)
    {
        rterror("init_yolo11_model fail! ret=%d model_path=%s\n", ret, model_path);
        return -1;
    }
    return 0;
}

int CMAIServiceImplRk3588::deinit_model() {
    deinit_post_process();

    int ret = release_yolo11_model(&rknn_app_ctx);
    if (ret != 0)
    {
        rterror("release_yolo11_model fail! ret=%d\n", ret);
    }
    return 0;
}

int CMAIServiceImplRk3588::recognize(const char *buf, size_t img_size, std::vector<CMAIRecogItem> &results) {
    image_buffer_t src_image;
    memset(&src_image, 0, sizeof(image_buffer_t));
    int ret = read_image_jpeg_buff(buf, img_size, &src_image);
    if (ret < 0)
    {
        rterror("read image failed [%d]\n", ret);
        return -2;
    }

    object_detect_result_list od_results;
    ret = inference_yolo11_model(&rknn_app_ctx, &src_image, &od_results);
    if (ret != 0)
    {
        rterror("init_yolo11_model fail! ret=%d\n", ret);
        return -2;
    }

    rtdebug("od_results count [%d]\n", od_results.count);
    // 画框和概率
    char text[256];
    for (int i = 0; i < od_results.count; i++)
    {
        object_detect_result *det_result = &(od_results.results[i]);
        rtdebug("%s @ (%d %d %d %d) %.3f\n", coco_cls_to_name(det_result->cls_id),
               det_result->box.left, det_result->box.top,
               det_result->box.right, det_result->box.bottom,
               det_result->prop);
        int x1 = det_result->box.left;
        int y1 = det_result->box.top;
        int x2 = det_result->box.right;
        int y2 = det_result->box.bottom;

//        draw_rectangle(&src_image, x1, y1, x2 - x1, y2 - y1, COLOR_BLUE, 3);

        sprintf(text, "(%d) %s %.1f%%", det_result->cls_id, coco_cls_to_name(det_result->cls_id), det_result->prop * 100);
        rtinfo("result: %s\n", text);
        std::string cls_name = coco_cls_to_name(det_result->cls_id);
        bool name_exist = false;
        for (auto &item : results)
        {
            if (item.name == cls_name)
            {
                name_exist = true;
            }
        }
        if (!name_exist)
        {
            results.push_back(CMAIRecogItem{coco_cls_to_name(det_result->cls_id),
                                           det_result->prop});
        }

//        draw_text(&src_image, text, x1, y1 - 20, COLOR_RED, 10);
    }
    if (src_image.virt_addr != NULL)
    {
#if defined(RV1106_1103)
        dma_buf_free(rknn_app_ctx.img_dma_buf.size, &rknn_app_ctx.img_dma_buf.dma_buf_fd,
                     rknn_app_ctx.img_dma_buf.dma_buf_virt_addr);
#else
        free(src_image.virt_addr);
#endif
    }
    return 0;
}
